from ultralytics import YOLO

# --- 1. 模型加载配置 ---
# 待转换的 YOLOv8 模型路径。
# 可以是本地训练好的 .pt 文件 (例如 'runs/detect/YOLOv8n_iflytekUcar/train/weights/best.pt')
# 也可以是 Ultralytics 提供的预训练模型名称 (例如 'yolov8n.pt')
MODEL_PATH = 'runs/detect/YOLOv8n_iflytekUcar/train/weights/best.pt' # 请替换为你的模型路径

print(f"正在加载模型: {MODEL_PATH} ...")
# 加载 YOLO 模型
model = YOLO(MODEL_PATH)
print(f"模型 '{MODEL_PATH}' 已成功加载。")

# --- 2. 模型导出参数配置 (高度可自定义) ---
# 这些参数将传递给 model.export() 方法，控制导出行为。
# 你可以根据需要调整这些值，以适应不同的部署需求和硬件平台。
EXPORT_PARAMS = {
    # 导出格式 (format): 指定模型转换的目标格式。
    # 常用选项包括:
    # 'onnx': 用于 ONNX Runtime 或后续转换为其他格式 (如 TensorRT, OpenVINO, RKNN)。
    # 'torchscript': PyTorch 自己的序列化格式。
    # 'openvino': 适用于 Intel CPU/GPU。
    # 'tensorrt': 适用于 NVIDIA GPU。
    # 'ncnn': 适用于 NCNN 框架。
    # 'coreml': 适用于 Apple 设备。
    # 'tflite': 适用于 TensorFlow Lite (移动/嵌入式设备)。
    # 'edgetpu': 适用于 Google Edge TPU。
    # 'rknn': 适用于瑞芯微 (Rockchip) 芯片。
    'format': 'rknn', # 默认为 'rknn'，请根据你的目标平台修改

    # 图像大小 (imgsz): 导出模型推理时期望的输入图像尺寸。
    # 可以是单个整数（如 640，表示 640x640），也可以是元组（如 (1920, 1080)）。
    # 确保此尺寸与你推理时使用的尺寸兼容。
    'imgsz': 640, # 默认 640。对于 ONNX/RKNN 等，通常固定输入尺寸。

    # 半精度导出 (half): 是否导出为 FP16 半精度模型。
    # 可以减小模型大小并加速推理，但可能会牺牲少量精度。
    # 某些硬件平台不支持 FP16。
    'half': False, # True 或 False

    # 设备 (device): 指定在导出过程中使用的设备。
    # 'cpu' 表示使用 CPU，0 表示使用第一个 GPU。
    # 如果导出复杂模型，建议使用 GPU 加速导出过程。
    'device': '0', # 'cpu' 或 '0'

    # 简化模型 (simplify): 是否使用 onnx-simplifier 工具简化 ONNX 模型。
    # 强烈推荐设置为 True，这对于后续转换为 RKNN 或其他嵌入式部署非常有用，能优化模型结构和性能。
    'simplify': True, # True 或 False (仅对 ONNX 格式有效)

    # NMS 处理 (nms): 是否在导出模型中包含非极大值抑制 (NMS) 逻辑。
    # 如果设置为 False，模型将输出原始的检测框和分数，NMS 步骤需要在推理端手动实现。
    # 这对于某些自定义部署流程（如在 RKNN 或其他专门硬件上实现NMS）非常重要，可以获得更灵活的控制。
    # 如果为 True，模型输出将是经过 NMS 处理的最终结果。
    'nms': False, # True 或 False。对于 RKNN 转换，通常设置为 False，然后在 RKNN SDK 中处理 NMS。

    # 推理批次大小 (batch): 导出模型支持的批次大小。
    # -1 表示动态批次大小（如果目标格式支持）。
    # 对于 ONNX，可以设置为 1 或 -1。对于大多数嵌入式设备，通常设置为 1。
    'batch': -1, # 例如 1 (固定批次为1) 或 -1 (动态批次)

    # 训练模式 (train): 是否以训练模式导出模型 (通常保持 False)。
    'train': False, # 默认 False

    # 卷积重参数化 (reparam): 是否在导出前对模型进行重参数化（例如 RepVGG）。
    # 通常保持默认，除非你明确知道需要。
    'reparam': False, # 默认 False

    # 输出名称 (name): 导出的模型文件名，不包含扩展名。
    # 默认情况下，会根据模型路径和格式自动生成，例如 'best.onnx'。
    # 如果设置此参数，它会作为基础文件名。
    # 'name': 'my_custom_exported_model', # 可以自定义文件名，例如 'my_custom_exported_model'
}

print("\n--- 开始导出模型 ---")
# 执行模型导出
# 导出后的文件路径将打印在控制台或位于 runs/export/ 目录下。
# 例如：'runs/detect/YOLOv8n_iflytekUcar/train/weights/best.onnx' 或 'runs/detect/YOLOv8n_iflytekUcar/train/weights/best.rknn'
exported_model_path = model.export(**EXPORT_PARAMS)

if exported_model_path:
    print(f"\n模型已成功导出为 {EXPORT_PARAMS['format'].upper()} 格式！")
    print(f"生成的文件位于: {exported_model_path}")
    print("\n下一步：")
    if EXPORT_PARAMS['format'] == 'onnx':
        print("1. 建议使用 Netron (https://netron.app/) 或 onnxsim 工具检查 ONNX 模型的结构。")
        print("2. 如果目标是 RKNN，现在可以进行 RKNN SDK 的转换步骤。")
    elif EXPORT_PARAMS['format'] == 'rknn':
        print("1. RKNN 模型已生成。请使用 RKNN Toolkit 进行推理测试或部署到 Rockchip 芯片。")
    else:
        print(f"请根据你选择的 {EXPORT_PARAMS['format'].upper()} 格式，进行相应的后续部署或验证。")
else:
    print("\n模型导出失败，请检查配置和错误信息。")